import sys
import threading
import Queue
import traceback


max_threads = 10
pool = None

class NoResultsPending(Exception):
    pass

class NoWorkersAvailable(Exception):
    pass

def _handle_thread_exception(request, exc_info):
    traceback.print_exception(*exc_info)
    
    
    
    
    
def make_requests(callable_, args_list, callback=None,
        exc_callback=_handle_thread_exception):

    requests = []
    for item in args_list:
        if isinstance(item, tuple):
            requests.append(
                WorkRequest(callable_, item[0], item[1], callback=callback,
                    exc_callback=exc_callback)
            )
        else:
            requests.append(
                WorkRequest(callable_, [item], None, callback=callback,
                    exc_callback=exc_callback)
            )
    return requests
    
    
    
    
class Worker( threading.Thread ):

    def __init__(self, in_queue, out_queue, poll_timeout=5, **kw):

        threading.Thread.__init__(self, **kw)
        
        self.setDaemon(1)
        self._in_queue = in_queue
        self._out_queue = out_queue
        self._poll_timeout = poll_timeout
        self._aborted = threading.Event()
        self.start()


    def run(self):
        while True:
            if self._aborted.isSet():
                break

            try:
                request = self._in_queue.get(True, self._poll_timeout)
            except Queue.Empty:
                continue
            else:
                if self._aborted.isSet():
                    self._out_queue.put(request)
                    break
                try:
                    result = request.callable(*request.args, **request.kw)
                    self._out_queue.put((request, result))
                except:
                    request.exception = True
                    self._out_queue.put((request, sys.exc_info()))


    def abort(self):
        self._aborted.set()
        
        
        


class WorkRequest:

    def __init__(self, callable_, args=None, kw=None, requestID=None,
            callback=None, exc_callback=_handle_thread_exception):

        if requestID is None:
            self.requestID = id(self)
        else:
            try:
                self.requestID = hash(requestID)
            except TypeError:
                raise TypeError("requestID must be hashable.")
        self.exception = False
        self.callback = callback
        self.exc_callback = exc_callback
        self.callable = callable_
        self.args = args or []
        self.kw = kw or {}

    def __str__(self):
        return "<WorkRequest id=%s args=%r kwargs=%r exception=%s>" % \
            (self.requestID, self.args, self.kwds, self.exception)
            
            
            
        
class Pool:

    def __init__(self, num_workers, q_size=0, resq_size=0, poll_timeout=5):
        self.workers = []
        self._in_queue = Queue.Queue(q_size)
        self._out_queue = Queue.Queue(resq_size)
        self.work_requests = {}
        self.aborted_workers = []
        self.create_workers( num_workers, poll_timeout )


    def create_workers(self, num_workers, poll_timeout=5):
        for i in range(num_workers):
            self.workers.append( Worker( self._in_queue,
                self._out_queue, poll_timeout=poll_timeout ))


    def abort_workers(self, num_workers, do_join=False):
        abort_list = []
        for i in range(min(num_workers, len(self.workers))):
            worker = self.workers.pop()
            worker.abort()
            abort_list.append(worker)

        if do_join:
            for worker in abort_list:
                worker.join()
        else: self.aborted_orkers.extend(abort_list)


    def join_all_aborted(self):
        for worker in self.aborted_orkers:
            worker.join()
        self.aborted_workers = []


    def add(self, request, block=True, timeout=None):
        assert isinstance(request, WorkRequest)
        assert not getattr(request, 'exception', None)
        self._in_queue.put(request, block, timeout)
        self.work_requests[request.requestID] = request


    def poll(self, block=False):
        while True:
            if not self.work_requests:
                raise NoResultsPending
            elif block and not self.workers:
                raise NoWorkersAvailable
            try:
                request, result = self._out_queue.get(block=block)
                if request.exception and request.exc_callback:
                    request.exc_callback(request, result)
                if request.callback and not \
                       (request.exception and request.exc_callback):
                    request.callback(request, result)
                del self.work_requests[request.requestID]
            except Queue.Empty:
                break


    def wait(self):
        while 1:
            try:
                self.poll(True)
            except NoResultsPending:
                break
                
                
                
                


def create_pool( threads=None ):
    global max_threads
    global pool
    if not pool:
        max_threads = threads or max_threads
        pool = Pool(max_threads)